import { IndexFlatL2, SearchResult } from "faiss-node";
import { promises as fs } from "fs";
import { resolve } from "path";
import { MemoryStore, MemoryStoreSaveArgs } from "./base";

/** 内存统计接口 */
interface MemoryStatistics {
    /** 总数据量 */
    size: number;
}

/** 存储内容及其嵌入的接口 */
interface StoredContent {
    id: number;
    embedding: number[];
    content: string;
}

/** 搜索结果对象 */
interface SearchResultObject {
    similarity: number;
    content: string;
}

/** JsonMemoryStore 实现 */
export class JsonMemoryStore extends MemoryStore<string> {
    private contents: StoredContent[] = [];
    private index: IndexFlatL2;
    private nextId: number = 0; // 用于分配唯一ID

    constructor(dimension: number) {
        super();
        this.index = new IndexFlatL2(dimension);
    }

    /**
     * 保存内容并添加到 Faiss 索引
     * @param content 要保存的内容
     * @param embedding 数据的 embedding
     * @returns 是否保存成功
     */
    async save(args: MemoryStoreSaveArgs<string>): Promise<boolean> {
        // 添加到内存存储
        const storedContent: StoredContent = {
            id: this.getNextId(),
            embedding: args.embedding,
            content: args.content
        };
        this.contents.push(storedContent);
        // 添加到 Faiss 索引
        this.index.add(args.embedding);
        return true;
    }

    /**
     * 删除具有特定嵌入的内容
     * @param contentEmbedding 要删除内容的嵌入
     * @returns 是否删除成功
     */
    async delete(contentEmbedding: number[]): Promise<boolean> {
        // 查找匹配的内容
        const index = this.contents.findIndex((item) =>
            this.compareEmbeddings(item.embedding, contentEmbedding)
        );

        if (index === -1) {
            return false; // 未找到匹配内容
        }

        // 移除内容
        const [removedContent] = this.contents.splice(index, 1);

        // 重新构建索引
        this.rebuildIndex();

        return true;
    }

    /**
     * 查询与提供的嵌入相似的内容
     * @param contentEmbedding 查询的嵌入
     * @param threshold 相似度阈值
     * @param maxCount 最大返回数量
     * @returns 相似内容及其相似度
     */
    async query(
        contentEmbedding: number[],
        threshold: number,
        maxCount: number
    ): Promise<SearchResultObject[]> {
        if (this.contents.length === 0) {
            return [];
        }
        const statistics = await this.getStatistics();
        if (statistics.size < maxCount) {
            maxCount = statistics.size;
        }

        if (contentEmbedding.length !== this.index.getDimension()) {
            throw new Error(
                `Query embedding dimension ${contentEmbedding.length} does not match index dimension ${this.index.getDimension()}`
            );
        }

        // 执行搜索
        const searchResult: SearchResult = this.index.search(contentEmbedding, maxCount);
        const results: SearchResultObject[] = [];

        for (let i = 0; i < searchResult.labels.length; i++) {
            const label = searchResult.labels[i];
            const distance = searchResult.distances[i];
            const similarity = 1 - distance / 2; // 余弦相似度近似

            if (similarity >= threshold && label !== -1) {
                const content = this.contents[label].content;
                results.push({ similarity, content });
            }
        }

        return results;
    }

    /**
     * 获取当前存储的统计数据
     * @returns 内存统计信息
     */
    async getStatistics(): Promise<MemoryStatistics> {
        return { size: this.contents.length };
    }

    /**
     * 重新构建 Faiss 索引
     */
    private rebuildIndex() {
        this.index = new IndexFlatL2(this.index.getDimension());
        const embeddings = this.contents.map((item) => item.embedding);
        if (embeddings.length > 0) {
            const flatEmbeddings = embeddings.flat();
            this.index.add(flatEmbeddings);
        }
    }

    /**
     * 比较两个嵌入是否相同
     * @param emb1 第一个嵌入
     * @param emb2 第二个嵌入
     * @returns 是否相同
     */
    private compareEmbeddings(emb1: number[], emb2: number[]): boolean {
        if (emb1.length !== emb2.length) return false;
        for (let i = 0; i < emb1.length; i++) {
            if (emb1[i] !== emb2[i]) return false;
        }
        return true;
    }

    /**
     * 将内存存储保存到 JSON 文件
     * @param filePath 文件路径
     */
    async saveToFile(filePath: string): Promise<void> {
        const path = resolve(filePath);
        const data = {
            nextId: this.nextId,
            contents: this.contents
        };
        await fs.writeFile(path, JSON.stringify(data, null, 2), "utf-8");
    }

    /**
     * 从 JSON 文件加载内存存储
     * @param filePath 文件路径
     */
    async loadFromFile(filePath: string): Promise<void> {
        const path = resolve(filePath);
        const data = await fs.readFile(path, "utf-8");
        const parsed = JSON.parse(data);
        this.nextId = parsed.nextId;
        this.contents = parsed.contents;

        // 重新构建索引
        this.rebuildIndex();
    }

    private getNextId() {
        return this.nextId++;
    }
}
